rm(list=ls())
library(rjags)
library(reshape)
root = "/Users/Tom/Documents/YNP_Bison_Model/"
#root = "/Volumes/wcnr-network/Research/Hobbs/A_Bison_2011/Analysis for Ecological Monograph/"

path = function (stem,r=root){
return(paste(r,stem, sep=""))	
}

setwd(path("A_final_model_runs_for_revised_manuscript/Density Dependent"))

set.seed(1)
#new data file
load(path("A_final_model_runs_for_revised_manuscript/2011 Bison Data with multinomial revised.Rdata"))


load(path("A_final_model_runs_for_revised_manuscript/Initial_state.Rdata"))

####set up dimensions of simulation
start.yr=1970
end.yr=2010
years=seq(start.yr,end.yr)
nyr=length(years)
pred.yr=0
n.class = 8 #number of classes in model
end.sim=nyr+pred.yr

#Function for making cells of projection matrix that hold parameters = NA
NA_matrix = function(nrow, ncol, nyr){
A=array(0,dim=c(nrow, ncol, nyr))
A[1,3,]<-NA
A[1,6,]<-NA
A[1,7,]<-NA

#row 2
A[2,1,]<-NA
#row 3
A[3,2, ]<-NA
A[3,3, ]<-NA

#row4
A[4,3, ]<-NA
A[4,6, ]<-NA
A[4,7, ]<-NA

#row5
A[5,1, ]<-NA
A[5,4, ]<-NA
#row6
A[6,2, ]<-NA
A[6,3, ]<-NA
A[6,5, ]<-NA
A[6,7, ]<-NA

#row 7
A[7,6, ]<-NA
A[7,7, ]<-NA

#row 8
A[8,1, ]<-NA
A[8,4, ]<-NA
A[8,8, ]<-NA
return(A)
}


library(rjags)
library(coda)

#add rows to Removal data to allow prediction
#pred.yr is the number of years to predict beyond 
#2008
#This asssumes 0 future removals.  Can be change in model experiments.
new.rN=matrix(0,nrow=pred.yr, ncol=n.class + 1)
new.yrs=seq(2009,2009+pred.yr-1)
new.rN[,1]=new.yrs

##make year index for non-zero removals
u=data$removals$count
remove.index=u[u$year >= 1985 & u$removal >0,1]


mydata = list(n.class = n.class,
	#y.nyr=nyr, 
	y.end.sim=end.sim,
	A = NA_matrix(nrow=n.class, ncol=n.class, nyr = end.sim),  #NA_matrix is a function that assins NA to elements of the matrix holding paramters
	#new.rN=new.rN[,2:(n.class +1)],
	#target = 0,
	#forecast.index=seq(nyr+1,end.sim),
		
	
	# Population composition data######################
	y.sigma.calf = data$aerial$calf$sd_of_mean,
	y.ratio.calf = data$aerial$calf$mean,
	y.iratio.calf = data$aerial$calf$ratio.index,
	
	y.ground.index=data$ground$mu[,2],
	y.p.mu = data$ground$alpha[,3:6]/rowSums(data$ground$alpha[,3:6]),
	#y.alpha = data$ground$alpha[,3:6],
	y.alpha_0=rowSums(data$ground$alpha[,3:6]),				
	#Count data########################################
	y.N=round(data$census$N$count.mean),
	y.N.var=data$census$N$count.sd^2,
	y.N.index=data$census$N$index, 	
	y.N.varindex = as.integer(data$census$sd_index$year.index), 
		#Seology data##################################
	#Calves
	y.pos.calf = data$sero$calf$boundary.calf.pos,
	y.n.calf = data$sero$calf$boundary.calf.tests,
	y.isero.calf = data$sero$calf$index,
	#Yearlings
	y.pos.yrl = data$sero$yrl$boundary.yrlng.pos,
	y.n.yrl = data$sero$yrl$boundary.yrlng.tests,
	y.isero.yrl = data$sero$yrl$index,
	#Adult cows
	y.pos.cow = data$sero$cow$boundary.adult.pos,
	y.n.cow = data$sero$cow$boundary.adult.tests,
	y.isero.cow = data$sero$cow$index,
	
	#Transmission data
	y.phi=data$trans_data$trans$mean,
	y.phi.sd=data$trans_data$trans$sd,
	y.iphi = data$trans_data$trans$index,

	
	#Transmission prior shape parameters
	y.shapes.v=data$trans_data$shapes.vert,
	y.shapes.psi=data$trans_data$shapes.psi,
	#y.shapes.gamma = data$trans_data$shapes.gamma,

	#removals
	y.rage=as.matrix(data$removals$comp[,3:6]),
	y.rsero.calf=as.matrix(data$removals$sero[,9:10]),
	y.rsero.yrlg=as.matrix(data$removals$sero[,11:12]),
	y.rsero.adult=as.matrix(data$removals$sero[,13:14]),
	y.removals=data$removals$count$removal,
	y.remove.index=remove.index,  
	y.I=diag(1,8)  
	)#end of data statement,
	
	
		
#These lines of code create an anlterative way to set initial values of the state vector. Convergence requires 8 x more iterations than when simulated, "ballpark", initial values are used, but the results are identical.	

#init_N=matrix(log(2000),ncol=n.class,nrow=end.sim)
#init_N[1,]=NA
############################
#initial values	
inits = list(
	list(
		psi = .01,
		b = c(.7,.4,.1),
		#alternatives for ininitializing state vector
	    #lambda=rep(2000,nyr),
		#log_N=init_N
		log_N=log(N.init[,,1]),
		lambda = round(rowSums(N.init[,,1])),
		m = .55,
		p=c(.75,.99,.80),
		beta = 1,
		#sigma.p=.1,
		sigma.p=.6,
		v = .4,
		r.sero.calf=rep(.20,end.sim),
		r.sero.yrlg=rep(.20,end.sim),
		r.sero.adult=rep(.20,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
		),
	
	
	list(
		psi = .07,
		b = c(.7,.6,.3),
		log_N=log(N.init[,,2]),
		lambda = round(rowSums(N.init[,,2])),
		# alternatives for ininitializing state vector
		# log_N=init_N/2,
		# lambda = rep(3000,nyr),
		#lambda = N.sim[1:nyr]*runif(nyr,.8,1.2),
		m = .6,
		p=c(.65,.90, .85),
		beta = .4,
		#tau.p = 5,
		sigma.p = .1,
		v = .2,
		r.sero.calf=rep(.10,end.sim),
		r.sero.yrlg=rep(.10,end.sim),
		r.sero.adult=rep(.10,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
	),

	list(
		psi = .03,
		b = c(.7,.6,.3),
		log_N=log(N.init[,,3]),
		lambda = round(rowSums(N.init[,,3])),
		#alternatives for ininitializing state vector
		#log_N=init_N*2,  
		#lambda = rep(4000,nyr),
		m = .6,
		p=c(.65,.90, .85),
		beta = 5,
		sigma.p = .15,
		v = .1,
		r.sero.calf=rep(.30,end.sim),
		r.sero.yrlg=rep(.30,end.sim),
		r.sero.adult=rep(.30,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
	))
	
beg.time=Sys.time()
n.iter=100000
n.update=50000
n.adapt=50000
model= "One_beta_model_finalJAGS.R"
jm=jags.model(model, data=mydata,inits,n.chains=length(inits), n.adapt = n.adapt)

beg.time=Sys.time()
update(jm,progress.bar="text", n.iter=n.update)
end.time=Sys.time()
run.time=end.time-beg.time

beg.time=Sys.time()
load.module("dic") #to get deviance in new version of JAGS. 
One_beta.coda=coda.samples(jm,variable.names=c("b", "p", "v", "psi", "tau.p", "sigma.p", "beta","m", "deviance", "N[35,3]"), n.iter=n.iter)
end.time=Sys.time()
run.time=end.time-beg.time

#summary(One_beta.coda)
#plot(One_beta.coda)

#gelman.diag(One_beta.coda)



One_beta.jags=jags.samples(jm,variable.names=c("N.total", "bull.ratio", "calf.ratio", "cow.ratio","yrl.ratio","pv.cow", "pv.calf","pv.yrl","phi", "in.per", "Re", "in.per.pos","vert.per","diff.pos","diff.con","b", "p", "v", "psi", "sigma.p", "beta","m", "deviance", "xcalf.ratio", "xcow.ratio", "xyrl.ratio", "xbull.ratio", "xpv.cow", "xpv.yrl", "xpv.calf", "xin.per", "xin.per.pos", "xvert.per", "xphi","xNewIP","mspe1", "mspe2", "P.sero.sd", "P.sero.sq", "P.sero.mean", "P.N.mean", "P.N.sd", "P.N.sq", "P.ratio.calf.mean", "P.ratio.calf.sd", "P.ratio.calf.sq" , "P_yphi_phi"), n.iter=n.iter)



save.image(file="One_beta_final.Rdata")
#save(One_beta.coda, file="One_beta_coda.Rdata")
#save(One_beta.jags, file="One_beta_jags.Rdata")